# Copyright (c) 2024, NVIDIA CORPORATION.

from enum import IntEnum

from pylibcudf.types import (
    DataType,
    Interpolation,
    NanEquality,
    NullEquality,
    NullOrder,
    NullPolicy,
    Order,
)

class Kind(IntEnum):
    SUM = ...
    PRODUCT = ...
    MIN = ...
    MAX = ...
    COUNT_VALID = ...
    COUNT_ALL = ...
    ANY = ...
    ALL = ...
    SUM_OF_SQUARES = ...
    MEAN = ...
    VARIANCE = ...
    STD = ...
    MEDIAN = ...
    QUANTILE = ...
    ARGMAX = ...
    ARGMIN = ...
    NUNIQUE = ...
    NTH_ELEMENT = ...
    RANK = ...
    COLLECT_LIST = ...
    COLLECT_SET = ...
    PTX = ...
    CUDA = ...
    CORRELATION = ...
    COVARIANCE = ...

class CorrelationType(IntEnum):
    PEARSON = ...
    KENDALL = ...
    SPEARMAN = ...

class EWMHistory(IntEnum):
    INFINITE = ...
    FINITE = ...

class RankMethod(IntEnum):
    FIRST = ...
    AVERAGE = ...
    MIN = ...
    MAX = ...
    DENSE = ...

class RankPercentage(IntEnum):
    NONE = ...
    ZERO_NORMALIZED = ...
    ONE_NORMALIZED = ...

class UdfType(IntEnum):
    CUDA = ...
    PTX = ...

class Aggregation:
    def __init__(self): ...
    def kind(self) -> Kind: ...

def sum() -> Aggregation: ...
def product() -> Aggregation: ...
def min() -> Aggregation: ...
def max() -> Aggregation: ...
def count(null_handling: NullPolicy = NullPolicy.INCLUDE) -> Aggregation: ...
def any() -> Aggregation: ...
def all() -> Aggregation: ...
def sum_of_squares() -> Aggregation: ...
def mean() -> Aggregation: ...
def variance(ddof: int = 1) -> Aggregation: ...
def std(ddof: int = 1) -> Aggregation: ...
def median() -> Aggregation: ...
def quantile(
    quantiles: list[float], interp: Interpolation = Interpolation.LINEAR
) -> Aggregation: ...
def argmax() -> Aggregation: ...
def argmin() -> Aggregation: ...
def ewma(center_of_mass: float, history: EWMHistory) -> Aggregation: ...
def nunique(null_handling: NullPolicy = NullPolicy.EXCLUDE) -> Aggregation: ...
def nth_element(
    n: int, null_handling: NullPolicy = NullPolicy.INCLUDE
) -> Aggregation: ...
def collect_list(
    null_handling: NullPolicy = NullPolicy.INCLUDE,
) -> Aggregation: ...
def collect_set(
    null_handling: NullPolicy = NullPolicy.INCLUDE,
    nulls_equal: NullEquality = NullEquality.EQUAL,
    nans_equal: NanEquality = NanEquality.ALL_EQUAL,
) -> Aggregation: ...
def udf(operation: str, output_type: DataType) -> Aggregation: ...
def correlation(type: CorrelationType, min_periods: int) -> Aggregation: ...
def covariance(min_periods: int, ddof: int) -> Aggregation: ...
def rank(
    method: RankMethod,
    column_order: Order = Order.ASCENDING,
    null_handling: NullPolicy = NullPolicy.EXCLUDE,
    null_precedence: NullOrder = NullOrder.AFTER,
    percentage: RankPercentage = RankPercentage.NONE,
) -> Aggregation: ...
def is_valid_aggregation(source: DataType, agg: Aggregation) -> bool: ...
